# 导入必要的模块和组件
from app.routers import openai_generate, openai_modify
from fastapi import FastAPI, Request, HTTPException, Depends, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware  # 导入CORS中间件，用于处理跨域资源共享
from slowapi import _rate_limit_exceeded_handler  # 导入速率限制异常处理器
from slowapi.errors import RateLimitExceeded  # 导入速率限制异常类
from app.routers import deepseek_generate, deepseek_modify  # 导入应用的路由模块
from app.core.limiter import limiter  # 导入速率限制器
from typing import cast, List, Optional, Dict, Any, Union
from pydantic import BaseModel  # 添加缺少的BaseModel导入
from starlette.exceptions import ExceptionMiddleware  # 导入异常中间件
from api_analytics.fastapi import Analytics  # 导入API分析中间件
import os  # 导入操作系统模块，用于获取环境变量
from dotenv import load_dotenv
import time
from app.routers import openai_doc, deepseek_doc  # 导入新的文档生成路由
from fastapi.responses import JSONResponse, StreamingResponse
import httpx
import openai
import json
import base64
import logging

# 加载环境变量
load_dotenv()

# 创建FastAPI应用实例
app = FastAPI(
    title="GitDiagram API",
    description="API for converting GitHub repositories to diagrams",
    version="0.1.0",
)

# 定义允许跨域请求的源列表
origins = ["http://localhost:3000", "https://gitdiagram.com", "http://103.133.177.116"]

# 添加CORS中间件，允许来自指定源的跨域请求
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,  # 允许的源列表
    allow_credentials=True,  # 允许携带凭证
    allow_methods=["GET", "POST"],  # 允许的HTTP方法
    allow_headers=["*"],  # 允许的HTTP头
)

# 从环境变量获取API分析密钥
API_ANALYTICS_KEY = os.getenv("API_ANALYTICS_KEY")
# 查看https://blog.csdn.net/gitblog_07673/article/details/142231051
if API_ANALYTICS_KEY:
    app.add_middleware(Analytics, api_key=API_ANALYTICS_KEY)

# 将速率限制器添加到应用状态
app.state.limiter = limiter
# 添加速率限制异常处理器
app.add_exception_handler(
    RateLimitExceeded, cast(ExceptionMiddleware, _rate_limit_exceeded_handler)
)

# 计算请求处理时间的中间件
@app.middleware("http")
async def add_process_time_header(request: Request, call_next):
    start_time = time.time()
    response = await call_next(request)
    process_time = time.time() - start_time
    response.headers["X-Process-Time"] = str(process_time)
    return response

# 健康检查接口
@app.get("/health")
async def health_check():
    return {"status": "ok"}

# 关于应用的接口
@app.get("/about")
async def about():
    return {
        "app": "GitDiagram API",
        "version": "0.1.0",
        "description": "API for converting GitHub repositories to diagrams and detailed documentation",
    }

# 注册路由模块
app.include_router(openai_generate.router)  # 注册OpenAI Mermaid图表生成路由
app.include_router(openai_modify.router)  # 注册修改图表的路由
app.include_router(deepseek_generate.router)  # 注册DeepSeek Mermaid图表生成路由
app.include_router(deepseek_modify.router)  # 注册DeepSeek修改图表的路由
app.include_router(openai_doc.router)       # 注册OpenAI生成文档的路由
app.include_router(deepseek_doc.router)     # 注册DeepSeek生成文档的路由

# 聊天消息模型
class ChatMessage(BaseModel):
    role: str
    content: str

# 聊天请求模型
class ChatRequest(BaseModel):
    messages: List[ChatMessage]
    api_key: Optional[str] = None
    model: str = "openai-o3"
    repository_url: Optional[str] = None
    code: Optional[str] = None
    document: Optional[str] = None

# 聊天响应模型
class ChatResponse(BaseModel):
    response: str

# OpenAI聊天接口
@app.post("/openai/chat", response_model=ChatResponse)
async def chat_with_openai(request: ChatRequest):
    if not request.api_key:
        raise HTTPException(status_code=400, detail="API key is required")
    
    try:
        openai_client = openai.OpenAI(api_key=request.api_key)
        
        # 准备消息，过滤并格式化
        messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
        
        # 添加仓库地址和代码上下文
        if request.repository_url or request.code or request.document:
            context = "Additional context:\n"
            if request.repository_url:
                context += f"Repository URL: {request.repository_url}\n\n"
            if request.code:
                context += f"Code snippet:\n```\n{request.code}\n```\n\n"
            if request.document:
                context += f"Document content:\n{request.document}\n\n"
            
            # 添加到系统消息
            for i, msg in enumerate(messages):
                if msg["role"] == "system":
                    messages[i]["content"] += f"\n\n{context}"
                    break
            else:
                # 如果没有系统消息，添加一个
                messages.insert(0, {
                    "role": "system",
                    "content": f"You are an AI assistant for the GitHub repository. {context}"
                })
        
        # 发送请求到OpenAI
        response = openai_client.chat.completions.create(
            model="gpt-3.5-turbo" if request.model == "openai-o3" else "gpt-4o-mini",
            messages=messages,
            temperature=0.7,
            max_tokens=2000
        )
        
        # 提取回复
        assistant_response = response.choices[0].message.content
        
        return ChatResponse(response=assistant_response)
    
    except Exception as e:
        logging.error(f"OpenAI chat error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error communicating with OpenAI: {str(e)}")

# OpenAI流式聊天接口
@app.post("/openai/chat/stream")
async def chat_with_openai_stream(request: ChatRequest):
    if not request.api_key:
        raise HTTPException(status_code=400, detail="API key is required")
    
    try:
        openai_client = openai.OpenAI(api_key=request.api_key)
        
        # 准备消息，过滤并格式化
        messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
        
        # 添加仓库地址和代码上下文
        if request.repository_url or request.code or request.document:
            context = "Additional context:\n"
            if request.repository_url:
                context += f"Repository URL: {request.repository_url}\n\n"
            if request.code:
                context += f"Code snippet:\n```\n{request.code}\n```\n\n"
            if request.document:
                context += f"Document content:\n{request.document}\n\n"
            
            # 添加到系统消息
            for i, msg in enumerate(messages):
                if msg["role"] == "system":
                    messages[i]["content"] += f"\n\n{context}"
                    break
            else:
                # 如果没有系统消息，添加一个
                messages.insert(0, {
                    "role": "system",
                    "content": f"You are an AI assistant for the GitHub repository. {context}"
                })
        
        # 流式响应生成器
        async def generate():
            # 发送请求到OpenAI
            stream = openai_client.chat.completions.create(
                model="gpt-3.5-turbo" if request.model == "openai-o3" else "gpt-4o-mini",
                messages=messages,
                temperature=0.7,
                max_tokens=2000,
                stream=True  # 启用流式响应
            )
            
            # 处理流式响应
            for chunk in stream:
                if chunk.choices[0].delta.content:
                    yield chunk.choices[0].delta.content
        
        # 返回流式响应
        return StreamingResponse(generate(), media_type="text/plain")
    
    except Exception as e:
        logging.error(f"OpenAI chat stream error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error communicating with OpenAI: {str(e)}")

# DeepSeek聊天接口
@app.post("/deepseek/chat", response_model=ChatResponse)
async def chat_with_deepseek(request: ChatRequest):
    if not request.api_key:
        raise HTTPException(status_code=400, detail="API key is required")
    
    try:
        # 准备消息，过滤并格式化
        messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
        
        # 添加仓库地址和代码上下文
        if request.repository_url or request.code or request.document:
            context = "Additional context:\n"
            if request.repository_url:
                context += f"Repository URL: {request.repository_url}\n\n"
            if request.code:
                context += f"Code snippet:\n```\n{request.code}\n```\n\n"
            if request.document:
                context += f"Document content:\n{request.document}\n\n"
            
            # 添加到系统消息
            for i, msg in enumerate(messages):
                if msg["role"] == "system":
                    messages[i]["content"] += f"\n\n{context}"
                    break
            else:
                # 如果没有系统消息，添加一个
                messages.insert(0, {
                    "role": "system",
                    "content": f"You are an AI assistant for the GitHub repository. {context}"
                })
        
        # 发送请求到DeepSeek API
        async with httpx.AsyncClient(timeout=60.0) as client:
            response = await client.post(
                "https://api.deepseek.com/v1/chat/completions",
                headers={
                    "Content-Type": "application/json",
                    "Authorization": f"Bearer {request.api_key}"
                },
                json={
                    "model": "deepseek-chat",
                    "messages": messages,
                    "temperature": 0.7,
                    "max_tokens": 2000
                }
            )
            
            if response.status_code != 200:
                raise HTTPException(
                    status_code=response.status_code,
                    detail=f"DeepSeek API error: {response.text}"
                )
            
            # 解析响应
            response_data = response.json()
            assistant_response = response_data["choices"][0]["message"]["content"]
            
            return ChatResponse(response=assistant_response)
    
    except Exception as e:
        logging.error(f"DeepSeek chat error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error communicating with DeepSeek: {str(e)}")

# DeepSeek流式聊天接口
@app.post("/deepseek/chat/stream")
async def chat_with_deepseek_stream(request: ChatRequest):
    if not request.api_key:
        raise HTTPException(status_code=400, detail="API key is required")
    
    try:
        # 准备消息，过滤并格式化
        messages = [{"role": msg.role, "content": msg.content} for msg in request.messages]
        
        # 添加仓库地址和代码上下文
        if request.repository_url or request.code or request.document:
            context = "Additional context:\n"
            if request.repository_url:
                context += f"Repository URL: {request.repository_url}\n\n"
            if request.code:
                context += f"Code snippet:\n```\n{request.code}\n```\n\n"
            if request.document:
                context += f"Document content:\n{request.document}\n\n"
            
            # 添加到系统消息
            for i, msg in enumerate(messages):
                if msg["role"] == "system":
                    messages[i]["content"] += f"\n\n{context}"
                    break
            else:
                # 如果没有系统消息，添加一个
                messages.insert(0, {
                    "role": "system",
                    "content": f"You are an AI assistant for the GitHub repository. {context}"
                })
        
        # 流式响应生成器
        async def generate():
            async with httpx.AsyncClient(timeout=60.0) as client:
                response = await client.post(
                    "https://api.deepseek.com/v1/chat/completions",
                    headers={
                        "Content-Type": "application/json",
                        "Authorization": f"Bearer {request.api_key}"
                    },
                    json={
                        "model": "deepseek-chat",
                        "messages": messages,
                        "temperature": 0.7,
                        "max_tokens": 2000,
                        "stream": True  # 启用流式响应
                    },
                    timeout=60.0
                )
                
                if response.status_code != 200:
                    raise HTTPException(
                        status_code=response.status_code,
                        detail=f"DeepSeek API error: {response.text}"
                    )
                
                # 解析流式响应
                async for line in response.aiter_lines():
                    line = line.strip()
                    if not line or line == "data: [DONE]":
                        continue
                    
                    # 移除SSE前缀'data: '
                    if line.startswith("data: "):
                        line = line[6:]
                    
                    try:
                        data = json.loads(line)
                        content = data.get("choices", [{}])[0].get("delta", {}).get("content", "")
                        if content:
                            yield content
                    except json.JSONDecodeError:
                        continue
        
        # 返回流式响应
        return StreamingResponse(generate(), media_type="text/plain")
    
    except Exception as e:
        logging.error(f"DeepSeek chat stream error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error communicating with DeepSeek: {str(e)}")

# 启动事件
@app.on_event("startup")
async def startup_event():
    print("Application startup")

# 关闭事件
@app.on_event("shutdown")
async def shutdown_event():
    print("Application shutdown")

# 定义根路由
@app.get("/")
# @limiter.limit("100/day")  # 速率限制装饰器（当前已注释掉）
async def root(request: Request):
    """根路由处理函数，返回简单的欢迎消息"""
    return {"message": "Hello from GitDiagram API!"}

@app.exception_handler(500)
async def internal_exception_handler(request: Request, exc: Exception):
    """处理内部服务器错误，返回友好的错误信息"""
    error_detail = str(exc) if app.debug else "Internal server error"
    print(f"[ERROR] 500 Internal Server Error: {str(exc)}")
    return JSONResponse(
        status_code=500,
        content={"error": error_detail, "status": "error"}
    )

@app.exception_handler(404)
async def not_found_exception_handler(request: Request, exc: Exception):
    """处理资源不存在错误"""
    return JSONResponse(
        status_code=404,
        content={"error": "The requested resource was not found", "status": "error"}
    )

@app.exception_handler(429)
async def rate_limit_exception_handler(request: Request, exc: Exception):
    """处理请求频率超限错误"""
    return JSONResponse(
        status_code=429,
        content={"error": "Too many requests. Please try again later.", "status": "error"}
    )
